# --------------------------------------------------------
# Pytorch Meta R-CNN
# Written by Anny Xu, Xiaopeng Yan, based on the code from Jianwei Yang
# --------------------------------------------------------
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.models as models
from torch.autograd import Variable
import numpy as np
from model.utils.config import cfg
from model.rpn.rpn import _RPN
#from model.roi_pooling.modules.roi_pool import _RoIPooling
#from model.roi_crop.modules.roi_crop import _RoICrop
#from model.roi_align.modules.roi_align import RoIAlignAvg
from model.roi_layers import ROIAlign, ROIPool
from model.rpn.proposal_target_layer_cascade import _ProposalTargetLayer
import time
import pdb
from model.utils.net_utils import _smooth_l1_loss, _crop_pool_layer, _affine_grid_gen, _affine_theta
import pickle


class _fasterRCNN(nn.Module):
    """ faster RCNN """

    def __init__(self, classes, class_agnostic, meta_train, meta_test=None, meta_loss=None, margin=None, sim=None, margin_value=0, gt_tsne=None, epoch=None, adaptive_margin=None):
        super(_fasterRCNN, self).__init__()
        self.classes = classes
        self.n_classes = len(classes)
        self.class_agnostic = class_agnostic
        self.meta_train = meta_train
        self.meta_test = meta_test
        self.meta_loss = meta_loss
        # loss
        self.RCNN_loss_cls = 0
        self.RCNN_loss_bbox = 0
        self.scale = 5
        self.output_size = 2
        self.sim = sim
        self.validation = None
        self.margin = margin
        self.margin_value = margin_value
        self.gt_tsne = gt_tsne
        self.epoch = epoch
        self.feature_nor = None
        self.adaptive_margin = adaptive_margin

        # define rpn
        self.RCNN_rpn = _RPN(self.dout_base_model)
        self.RCNN_proposal_target = _ProposalTargetLayer(self.n_classes)
        self.RCNN_roi_pool = ROIPool((cfg.POOLING_SIZE, cfg.POOLING_SIZE), 1.0 / 16.0)
        self.RCNN_roi_align = ROIAlign((cfg.POOLING_SIZE, cfg.POOLING_SIZE), 1.0 / 16.0, 0)

        self.grid_size = cfg.POOLING_SIZE * 2 if cfg.CROP_RESIZE_WITH_MAX_POOL else cfg.POOLING_SIZE
        self.num_rois = torch.zeros(21, 4) # num, ratio_sum, count, fail_count
        self.iou_list = [[] for _ in range(21)]
        self.au = False
        

    def forward(self, im_data_list, im_info_list, gt_boxes_list, num_boxes_list, average_shot=None,
                mean_class_attentions=None):
        # return attentions for testing
        if average_shot:
            prn_data = im_data_list[0]  # len(metaclass)*4*224*224
            attentions = self.prn_network(prn_data)
            return attentions  
        # extract attentions for training
        if self.meta_train and self.training and not self.validation:
            prn_data = im_data_list[0]  # len(metaclass)*4*224*224
            # feed prn data to prn_network
            attentions = self.prn_network(prn_data)
            prn_cls = im_info_list[0]  # len(metaclass)

        im_data = im_data_list[-1]
        im_info = im_info_list[-1]
        gt_boxes = gt_boxes_list[-1]
        num_boxes = num_boxes_list[-1]

        batch_size = im_data.size(0)
        im_info = im_info.data
        gt_boxes = gt_boxes.data
        num_boxes = num_boxes.data

        # feed image data to base model to obtain base feature map
        base_feat = self.RCNN_base(self.rcnn_conv1(im_data))

        if self.meta_train and self.epoch > 20 and self.adaptive_margin:
            # pdb.set_trace()
            gt_boxes_1 = torch.zeros(gt_boxes_list[0].data.shape[0] * gt_boxes_list[0].data.shape[1], 5)
            gt_boxes_1[:, 1:] = gt_boxes_list[0].data.view(-1, 5)[:, :4]
            true_label = gt_boxes_list[0].data.view(-1, 5)[:, 4].cpu()
            pooled_feat = self.RCNN_roi_pool(base_feat, Variable(gt_boxes_1).cuda()).detach()

            pooled_feat = self._head_to_tail(pooled_feat).detach()  # n*2048
            pooled_feat = np.array(pooled_feat.data.cpu())
            true_label = np.array(true_label)
            keep = true_label > 0
            pooled_feat = pooled_feat[keep]
            true_label = true_label[keep]
            with open('/home/wjh/Meta_R-CNN/logs/gt_box_feature_%.0f.txt' % self.epoch, 'a') as f:
                    f.write(str(np.array(pooled_feat).tolist()) + '\n')
                    # pdb.set_trace()

            with open('/home/wjh/Meta_R-CNN/logs/true_label_%.0f.txt' % self.epoch, 'a') as f:
                    f.write(str(np.array(true_label).tolist()) + '\n')

        # test with gt boxes
        if self.gt_tsne: 
            cls_prob_list = []
            channel_wise_feat_list = []
            gt_boxes = torch.zeros(gt_boxes_list[0].data.shape[0] * gt_boxes_list[0].data.shape[1], 5)
            gt_boxes[:, 1:] = gt_boxes_list[0].data.view(-1, 5)[:, :4]
            true_label = gt_boxes_list[0].data.view(-1, 5)[:, 4].cpu()
            pooled_feat = self.RCNN_roi_pool(base_feat, Variable(gt_boxes).cuda())
            
            pooled_feat = self._head_to_tail(pooled_feat)  # (b*128)*2048
            
            for i in range(len(mean_class_attentions)):
                mean_attentions = mean_class_attentions[i]
                channel_wise_feat = pooled_feat * mean_attentions
    
                # compute object classification probability
                cls_score = self.RCNN_cls_score(channel_wise_feat)
                        
                if self.margin:
                    cls_score -= cls_score.min(dim=1, keepdim=True)[0]
                    cls_score = self.scale * cls_score
                
                cls_prob = F.softmax(cls_score, dim = 1)

                cls_prob = cls_prob.view(batch_size, gt_boxes.size(0), -1)
                cls_prob_list.append(cls_prob)
                channel_wise_feat_list.append(channel_wise_feat.cpu())

            return true_label, cls_prob_list, channel_wise_feat_list
        
        # feed base feature map tp RPN to obtain rois
        rois, rpn_loss_cls, rpn_loss_bbox = self.RCNN_rpn(base_feat, im_info, gt_boxes, num_boxes)

        # if it is training phase, then use ground truth bboxes for refining
        if self.training and not self.validation:
            roi_data = self.RCNN_proposal_target(rois, gt_boxes, num_boxes)
            rois, rois_label, rois_target, rois_inside_ws, rois_outside_ws = roi_data
            rois_label = Variable(rois_label.view(-1).long()).detach()
            rois_target = Variable(rois_target.view(-1, rois_target.size(2))).detach()
            rois_inside_ws = Variable(rois_inside_ws.view(-1, rois_inside_ws.size(2))).detach()
            rois_outside_ws = Variable(rois_outside_ws.view(-1, rois_outside_ws.size(2))).detach()
        else:
            roi_data = self.RCNN_proposal_target(rois, gt_boxes, num_boxes)
            rois, rois_label, rois_target, rois_inside_ws, rois_outside_ws = roi_data
            
            
            rois_bboxes = rois.squeeze(0)[:, 1:5]
            gt_boxes = gt_boxes.squeeze(0)
            for gt_box in gt_boxes:
                gt_cls = int(gt_box[4])
                gt_box = gt_box.repeat(rois_bboxes.shape[0], 1)
                
                s_1 = (rois_bboxes[:, 2] - rois_bboxes[:, 0]) * (rois_bboxes[:, 3] - rois_bboxes[:, 1])
                s_2 = (gt_box[:, 2] - gt_box[:, 0]) * (gt_box[:, 3] - gt_box[:, 1])
                
                s_sum = s_1 + s_2
                
                left_line = torch.max(rois_bboxes[:, 1], gt_box[:, 1])
                right_line = torch.min(rois_bboxes[:, 3], gt_box[:, 3])
                top_line = torch.max(rois_bboxes[:, 0], gt_box[:, 0])
                bottom_line = torch.min(rois_bboxes[:, 2], gt_box[:, 2])
                
                s_inter = (right_line - left_line) * (bottom_line - top_line) * ((right_line - left_line) > 0).float() * ((bottom_line - top_line) > 0).float()
                iou_value = s_inter / (s_sum - s_inter)
                
                
                if torch.sum(iou_value > 0.5) == 0: # can not find any recall of this RoI
                    self.num_rois[gt_cls, 3] = self.num_rois[gt_cls, 3] + 1
                    num = 0
                    recall = 0
                else:
                    num = iou_value[iou_value > 0.5].size()[0]
                    recall = iou_value[iou_value > 0.5].size()[0] / rois_bboxes.shape[0] # the recall of RPN
                
                self.num_rois[gt_cls, 0] = self.num_rois[gt_cls, 0] + num
                self.num_rois[gt_cls, 1] = self.num_rois[gt_cls, 1] + recall
                self.num_rois[gt_cls, 2] = self.num_rois[gt_cls, 2] + 1
                self.iou_list[gt_cls] = self.iou_list[gt_cls] + iou_value[iou_value > 0.5].tolist()
            
            # rois_label = None
            rois_target = None
            rois_inside_ws = None
            rois_outside_ws = None
            rpn_loss_cls = 0
            rpn_loss_bbox = 0
          
        rois = Variable(rois)
        # do roi pooling based on predicted rois
        if cfg.POOLING_MODE == 'align':
            pooled_feat = self.RCNN_roi_align(base_feat, rois.view(-1, 5))  # (b*128)*1024*7*7
        elif cfg.POOLING_MODE == 'pool':
            pooled_feat = self.RCNN_roi_pool(base_feat, rois.view(-1, 5))
            
        pooled_feat = self._head_to_tail(pooled_feat)  # (b*128)*2048
        
        # meta training phase
        if self.meta_train and not self.validation:
            rcnn_loss_cls = []
            rcnn_loss_bbox = []
            if self.au:
                rcnn_loss_cls_base = []
                rcnn_loss_cls_novel = []
                rcnn_loss_bbox_base = []
                rcnn_loss_bbox_novel = []
            zero = Variable(torch.FloatTensor([0]).cuda())
            # pooled feature maps need to operate channel-wise multiplication with the corresponding class's attentions of every roi of image
            
            for b in range(batch_size):
                proposal_labels = rois_label[b * 128:(b + 1) * 128].data.cpu().numpy()
                unique_labels = list(np.unique(proposal_labels)) # the unique rois labels of the input image
                
                for i in range(attentions.size(0)):  # attentions len(attentions)*2048
                    if prn_cls[i].numpy()[0] + 1 not in unique_labels:
                        rcnn_loss_cls.append(zero)
                        rcnn_loss_bbox.append(zero)
                        continue

                    channel_wise_feat = pooled_feat[b * cfg.TRAIN.BATCH_SIZE:(b + 1) * cfg.TRAIN.BATCH_SIZE, :] * \
                                            attentions[i]  # 128x2048 channel-wise multiple
                    bbox_pred = self.RCNN_bbox_pred(channel_wise_feat)  # 128 * 4
                    # pdb.set_trace()
                    if self.training and not self.class_agnostic:
                        # select the corresponding columns according to roi labels
                        bbox_pred_view = bbox_pred.view(bbox_pred.size(0), int(bbox_pred.size(1) / 4), 4)
                        bbox_pred_select = torch.gather(bbox_pred_view, 1, rois_label[b * cfg.TRAIN.BATCH_SIZE:(b + 1) * cfg.TRAIN.BATCH_SIZE].view(rois_label[b * cfg.TRAIN.BATCH_SIZE:(b + 1) * cfg.TRAIN.BATCH_SIZE].size(0), 1, 1).expand(rois_label[b * cfg.TRAIN.BATCH_SIZE:(b + 1) * cfg.TRAIN.BATCH_SIZE].size(0), 1, 4))
                        bbox_pred = bbox_pred_select.squeeze(1)
                                            
                    # compute object classification probability
                    cls_score = self.RCNN_cls_score(channel_wise_feat)  # 128 * 21       

                    if self.training:
                        # classification loss
                        
                        if self.au:
                            keep_base = rois_label[b * 128:(b + 1) * 128] < 16
                            keep_novel = rois_label[b * 128:(b + 1) * 128] >= 16
                            keep_base_tmp = torch.nonzero(keep_base).squeeze()
                            keep_novel_tmp = torch.nonzero(keep_novel).squeeze()
                            if torch.sum(keep_base).data.cpu().numpy()[0] == 0:
                                RCNN_loss_cls_base = zero
                            else:
                                RCNN_loss_cls_base = F.cross_entropy(cls_score[keep_base_tmp, :], rois_label[b * 128:(b + 1) * 128][keep_base_tmp])
                            if torch.sum(keep_novel).data.cpu().numpy()[0] == 0:
                                RCNN_loss_cls_novel = zero
                            else:
                                RCNN_loss_cls_novel = F.cross_entropy(cls_score[keep_novel_tmp, :], rois_label[b * 128:(b + 1) * 128][keep_novel_tmp])
                            rcnn_loss_cls_base.append(RCNN_loss_cls_base)
                            rcnn_loss_cls_novel.append(RCNN_loss_cls_novel)
                        else:
                            RCNN_loss_cls = F.cross_entropy(cls_score, rois_label[b * 128:(b + 1) * 128])
                            rcnn_loss_cls.append(RCNN_loss_cls)
                        
                        # bounding box regression L1 loss
                        RCNN_loss_bbox = _smooth_l1_loss(bbox_pred, rois_target[b * 128:(b + 1) * 128],
                                                                rois_inside_ws[b * 128:(b + 1) * 128],
                                                                rois_outside_ws[b * 128:(b + 1) * 128])
                            
                        rcnn_loss_bbox.append(RCNN_loss_bbox)
            
            # meta attentions loss
            if self.meta_loss:
                attentions_score = self.Meta_cls_score(attentions)
                if self.au:
                    meta_loss_base = F.cross_entropy(attentions_score[:-5, :], Variable(torch.cat(prn_cls,dim=0).cuda())[:-5])
                    meta_loss_novel = F.cross_entropy(attentions_score[-5:, :], Variable(torch.cat(prn_cls,dim=0).cuda())[-5:])
                else:
                    meta_loss = F.cross_entropy(attentions_score, Variable(torch.cat(prn_cls,dim=0).cuda()))
            else:
                meta_loss = 0
                meta_loss_base = 0
                meta_loss_novel = 0
            
            if self.au:
                return rpn_loss_cls, rpn_loss_bbox, rcnn_loss_cls, rcnn_loss_bbox, rois_label, meta_loss_base, meta_loss_novel, rcnn_loss_cls_base, rcnn_loss_cls_novel, rcnn_loss_bbox_base, rcnn_loss_bbox_novel
            else:
                return rpn_loss_cls, rpn_loss_bbox, rcnn_loss_cls, rcnn_loss_bbox, rois_label, meta_loss

        elif self.meta_test:
            cls_prob_list = []
            bbox_pred_list = []
            sim_cls_fg_list = []
            gt_cls_fg_list = []
            gt_cls_all_list = []
            pre_cls_list = []
            reweighted_pre_cls_list = []
            channel_wise_feat_list = []

            for i in range(len(mean_class_attentions)):
                mean_attentions = mean_class_attentions[i]
                channel_wise_feat = pooled_feat * mean_attentions
                # compute bbox offset
                bbox_pred = self.RCNN_bbox_pred(channel_wise_feat)
                if self.training and not self.class_agnostic:
                    # select the corresponding columns according to roi labels
                    bbox_pred_view = bbox_pred.view(bbox_pred.size(0), int(bbox_pred.size(1) / 4), 4)
                    bbox_pred_select = torch.gather(bbox_pred_view, 1,
                                                    rois_label.view(rois_label.size(0), 1, 1).expand(rois_label.size(0),
                                                                                                        1, 4))
                    bbox_pred = bbox_pred_select.squeeze(1)

                # compute object classification probability
                cls_score = self.RCNN_cls_score(channel_wise_feat)
                    
                if self.margin:
                    cls_score -= cls_score.min(dim=1, keepdim=True)[0]
                    cls_score = self.scale * cls_score
                
                cls_prob = F.softmax(cls_score, dim = 1)

                RCNN_loss_cls = 0
                RCNN_loss_bbox = 0

                if self.training:
                    # classification loss
                    RCNN_loss_cls = F.cross_entropy(cls_score, rois_label)
                    # bounding box regression L1 loss
                    RCNN_loss_bbox = _smooth_l1_loss(bbox_pred, rois_target, rois_inside_ws, rois_outside_ws)

                cls_prob = cls_prob.view(batch_size, rois.size(1), -1)
                bbox_pred = bbox_pred.view(batch_size, rois.size(1), -1)
                cls_prob_list.append(cls_prob)
                bbox_pred_list.append(bbox_pred)
                channel_wise_feat_list.append(channel_wise_feat)

            return rois, rpn_loss_cls, rpn_loss_bbox, RCNN_loss_cls, RCNN_loss_bbox, rois_label, cls_prob_list, bbox_pred_list, 0, 0, 0, sim_cls_fg_list, \
                    gt_cls_fg_list, gt_cls_all_list, pre_cls_list, reweighted_pre_cls_list, channel_wise_feat_list
        else:
            bbox_pred = self.RCNN_bbox_pred(pooled_feat)
            if self.training and not self.class_agnostic:
                # select the corresponding columns according to roi labels
                bbox_pred_view = bbox_pred.view(bbox_pred.size(0), int(bbox_pred.size(1) / 4), 4)
                bbox_pred_select = torch.gather(bbox_pred_view, 1,
                                                rois_label.view(rois_label.size(0), 1, 1).expand(rois_label.size(0), 1,
                                                                                                 4))
                bbox_pred = bbox_pred_select.squeeze(1)

            cls_score = self.RCNN_cls_score(pooled_feat)
            cls_prob = F.softmax(cls_score, dim = 1)

            RCNN_loss_cls = 0
            RCNN_loss_bbox = 0

            if self.training:
                # classification loss
                RCNN_loss_cls = F.cross_entropy(cls_score, rois_label)
                # bounding box regression L1 loss
                RCNN_loss_bbox = _smooth_l1_loss(bbox_pred, rois_target, rois_inside_ws, rois_outside_ws)

            cls_prob = cls_prob.view(batch_size, rois.size(1), -1)
            bbox_pred = bbox_pred.view(batch_size, rois.size(1), -1)

        return rois, rpn_loss_cls, rpn_loss_bbox, RCNN_loss_cls, RCNN_loss_bbox, rois_label, cls_prob, bbox_pred, 0, 0, 0, None, None, None, None, None, pooled_feat

    def _init_weights(self):
        def normal_init(m, mean, stddev, truncated=False):
            """
            weight initalizer: truncated normal and random normal.
            """
            # x is a parameter
            if truncated:
                m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean)  # not a perfect approximation
            else:
                m.weight.data.normal_(mean, stddev)
                m.bias.data.zero_()

        normal_init(self.RCNN_rpn.RPN_Conv, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_rpn.RPN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_rpn.RPN_bbox_pred, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.RCNN_bbox_pred, 0, 0.001, cfg.TRAIN.TRUNCATED)

    def create_architecture(self):
        self._init_modules()
        self._init_weights()
